#!/usr/bin/env python3
# vim:fileencoding=utf-8

import os
import sys
import base64
import quopri
from mailbox import mbox
import email.header

from lxml.html import fromstring, tostring
from lxml.html import builder as E
from termcolor import colored

from charset import CJK_align
from myutils import filesize

mailStyle = '''\
    body { max-width: 900px; margin: auto; padding: 1em; }
    blockquote { margin-top: 1em !important; }
'''

def decodeHeader(h):
  var = email.header.decode_header(h)[0]
  charset = var[1] or 'ascii'
  if charset.lower() == 'gb2312': #fxxk
    charset = 'gb18030'
  try:
    var = var[0].decode(charset)
  except AttributeError:
    var = var[0]
  return var

def getMailSize(mail):
  payload = mail.get_payload()
  if isinstance(payload, str):
    return len(payload)
  else:
    return sum(getMailSize(m) for m in payload)

def selectMail(mb):
  for i, m in enumerate(mb):
    who = m['From']
    if who.endswith('>'):
      who = who.split(' <', 1)[0].strip('"')
      who = decodeHeader(who)
    else:
      who = who.split('@', 1)[0]
    subj = decodeHeader(m['Subject'])
    size = getMailSize(m)
    print(colored('%3d' % i, attrs=['bold']),
          #FIXME: strip overflowed text
          colored(CJK_align(who, 12), 'blue'),
          colored('(%7s)' % filesize(size)[:-1].rstrip('i'), 'yellow'),
          subj)

  while True:
    try:
      n = int(input('选择一封邮件: '))
      if n < 0 or n >= len(mb):
        raise ValueError
      return n
    except ValueError:
      continue
    except (EOFError, KeyboardInterrupt):
      print()
      raise SystemExit

def parseSingleMail(mail):
  mainMail = [m for m in mail.get_payload() if m.get_content_type() == 'text/html'][0]
  mailbody = getMailContent(mainMail)
  return mailbody

def saveHTMLMail(m):
  title = decodeHeader(m['Subject'])
  mailtype = m.get_content_type()
  if mailtype == 'multipart/alternative':
    mailbody = parseSingleMail(m)
  elif mailtype in ('multipart/related', 'multipart/mixed'):
    mails = m.get_payload()
    cidMapping = {}
    for mail in mails:
      if mail.get_content_type() == 'multipart/alternative':
        mailbody = parseSingleMail(mail)
      else:
        try:
          cid = mail['Content-ID'][1:-1]
        except TypeError:
          if mail['Content-Disposition'] and \
             mail['Content-Disposition'].find('attachment') != -1:
            continue
          from cli import repl
          repl(locals())
        fname = decodeHeader(mail.get_filename())
        cidMapping[cid] = fname
        body = getMailContent(mail)
        saveFile(fname, body)
  elif mailtype == 'text/plain':
    print('plain text mail, nothing to do')
    return
  elif mailtype == 'text/html':
    mailbody = getMailContent(m)
  else:
    raise NotImplementedError('type %s not recognized' % mailtype)

  div = fromstring(mailbody)
  for cidLink in div.cssselect('[src^="cid:"]'):
    cid = cidLink.get('src')[4:]
    cidLink.set('src', cidMapping[cid])
  div.insert(0, E.TITLE(title))
  div.insert(0, E.STYLE(mailStyle))
  mailbody_b = tostring(div, encoding='utf-8')
  saveFile('index.html', mailbody_b)

def saveFile(fname, content):
  if isinstance(content, str):
    f = open(fname, 'w')
  else:
    f = open(fname, 'wb')
  f.write(content)

def getMailContent(mail):
  rawbody = mail.get_payload()
  encoding = mail['Content-Transfer-Encoding']
  if encoding == 'base64':
    mailbody = base64.decodebytes(rawbody.encode('ascii'))
  elif encoding == 'quoted-printable':
    mailbody = quopri.decodestring(rawbody.encode('ascii'))
  else:
    raise NotImplementedError('encoding %s not recognized' % encoding)

  charset = mail.get_content_charset()
  if charset:
    mailbody = mailbody.decode(charset)
  return mailbody

def main(mailbox=os.path.expanduser('~/.Mail/inbox')):
  mb = mbox(mailbox)
  n = selectMail(mb)
  saveHTMLMail(mb[n])

if __name__ == '__main__':
  if len(sys.argv) == 2:
    main(sys.argv[1])
  elif len(sys.argv) == 1:
    main()
  else:
    print('usage: %s [MBOX_FILE]' % os.path.split(sys.argv[0])[1],
          file=sys.stderr)
